/*
 *  rootlocuswidget.cpp
 *  Perceel
 *
 *  Created by Nicola Gigante on 30/03/08.
 *  Copyright 2008 Nicola Gigante. All rights reserved.
 *
 */

#include "rootlocuswidget.h"

RootLocusWidget::RootLocusWidget(QWidget *parent) : PlotWidget(parent), m_convergence(true)
{
	
}

bool RootLocusWidget::checkPhase(Complex s)
{
	long double p = fabs(arg(tf().eval(s)));
	const long double p1 = M_PI;
	const long double p2 = M_PI + M_PI*2;
	const long double p3 = M_PI + M_PI*4;
	const long double k = 200;
	
	return (p > p1 - p1/k && p < p1 + p1/k) || (p > p2 - p1/k && p < p2 + p1/k) || (p > p3 - p1/k && p < p3 + p1/k);
}

QMap<QString, QString> RootLocusWidget::pixelDetails(QPoint p)
{
	QMap<QString, QString> m;
	
	Complex s = unMap(p);
	if(checkPhase(s))
	{
		Real k = real(-pow(tf().eval(s), -1));
		m[tr("K")] = QString::number(k);
	}else
		m[tr("K")] = tr("NaN");
	return m;
}

void RootLocusWidget::initPlot(QPainter &)
{
	setOrigin(QPointF(size().width() / 2, size().height() / 2));
}

void RootLocusWidget::plot(QPainter &painter)
{	
	painter.setPen(QPen(Qt::green, 3));
	painter.setBrush(Qt::green);
	
	if(m_convergence)
	{
		int iters = 0;
		fillLocusAxis(painter);
		QPointF prevPosPoint, prevNegPoint;
		bool firstPos = true, firstNeg = true;
		for(Real K = 0; K < 300; K += 0.1)
		{
			Polynomial poly = tf().denominator() + tf().numerator()*K;
			QList<Complex> roots;
		
			try
			{
				roots = poly.roots();
			}catch(QString &)
			{
				K = 0;
				if(++iters == 100)
				{
					convergenceError(painter);
					m_convergence = false;
					break;
				}
				continue;
			}
			
			for(int i = 0; i < roots.count(); i++)
			{
				if(isReal(roots.at(i)))
					continue;
				
				if(imag(roots.at(i)) > 0)
				{
					QPointF p = map(roots.at(i));
					if(firstPos)
					{
						prevPosPoint = p;
						firstPos = false;
					}
					painter.drawLine(prevPosPoint, p);
					prevPosPoint = p;
				}
				else
				{
					QPointF p = map(roots.at(i));
					if(firstNeg)
					{
						prevNegPoint = p;
						firstNeg = false;
					}
					painter.drawLine(prevNegPoint, p);
					prevNegPoint = p;
				}
			}
		}
		drawMarks(painter);
	}else
		convergenceError(painter);
	
}

void RootLocusWidget::fillLocusAxis(QPainter &painter)
{
	QList<Complex> l = tf().roots() + tf().poles();
	qSort(l.begin(), l.end());
	
	QPointF p1(rect().topLeft().x(), map(QPointF(verticalAxis(), horizontalAxis())).y());
	QPointF p2;
	Complex startC = unMap(p1);
	
	for(int i = 0; i < l.count(); i++)
	{
		if(real(l.at(i)) < real(startC) || !isReal(l.at(i)))
			continue;
		p2 = map(l.at(i));
		if((l.count() - i) % 2 != 0)
			painter.drawLine(p1, p2);
		p1 = p2;
	}
}

void RootLocusWidget::guessScale()
{
	Real max = 1;
	
	for(int i = 0; i < tf().roots().count(); i++)
	{
		Real r = abs(tf().roots().at(i));
		max = max > r ? max : r;
	}
	
	for(int i = 0; i < tf().poles().count(); i++)
	{
		Real p = abs(tf().poles().at(i));
		max = max > p ? max : p;
	}
	
	max *= 2.1;
	
	setScale(width() / max, height() / max);
}

void RootLocusWidget::drawMarks(QPainter &painter)
{
	QList<Complex> roots = tf().roots();
	QList<Complex> poles = tf().poles();
	
	painter.save();
	QPen p = painter.pen();
	p.setColor(Qt::blue);
	painter.setPen(p); // Roots color
	QBrush b = painter.brush();
	b.setColor(Qt::blue);
	painter.setBrush(b);
	for(int i = 0; i < roots.count(); i++)
	{
		QPointF p = map(roots.at(i));
		painter.drawEllipse(p.x() - 5, p.y() - 5, 10, 10);
		painter.drawText(p.x(), p.y() + painter.fontMetrics().height(), complexToString(roots.at(i)));
	}
	
	p.setColor(Qt::red);
	painter.setPen(p); // Poles color
	for(int i = 0; i < poles.count(); i++)
	{
		QPointF p = map(poles.at(i));
		QRect r = QRect(p.x() - 5, p.y() - 5, 10, 10);
		painter.drawLine(r.topLeft(), r.bottomRight());
		painter.drawLine(r.topRight(), r.bottomLeft());
		painter.drawText(p.x(), p.y() + painter.fontMetrics().height(), complexToString(poles.at(i)));
	}
	painter.restore();
}

void RootLocusWidget::convergenceError(QPainter &painter)
{
	QString error = tr("The root-finding algorithm could not converge.\n"
					   "Unable to plot the root locus of this transfer function.\n"
					   "We'll address this issue in a future release.");
	
	painter.save();
	painter.setPen(Qt::red);
	painter.fillRect(rect(), Qt::black);
	
	QStringList l = error.split('\n');
	int y = height() / 2;
	
	for(int i = 0; i < l.count(); i++)
	{
		QSize s = painter.fontMetrics().boundingRect(l.at(i)).size();
		QPoint p = QPoint((width() - s.width()) / 2, y);
		painter.drawText(p, l.at(i));
		y += s.height();
	}
	
	painter.restore();
}
